import os
import pymupdf
from rapidocr_onnxruntime import RapidOCR
import tempfile
import json
config = json.load(open("config.json"))
# 创建文档
outdoc = pymupdf.open()

# 加载外部字体（假设字体文件为 "myfont.ttf"）
font = pymupdf.Font(fontname="wqyyh", fontfile="fonts/"+config["def_font"])


def read_img(img_path) -> str:
    """使用RapidOCR识别图片的文字内容"""
    engine = RapidOCR(rec_model_path="ocrmodel/ch_PP-OCRv4_rec_infer.onnx")
    result, elapse = engine(img_path)
    resText = []
    for res in result:
        resText.append(res[1])
    return "\n".join(resText)


def extract_text_from_pdf(pdf_path, output_file, max_page=10):
    """从pdf读取文本，并通过rapidocr识别图片中的文字"""

    # 打开PDF文件
    doc = pymupdf.open(pdf_path)
    # 初始化一个空字符串来存储提取的文字
    # 遍历每一页
    for page_num in range(max(len(doc), max_page)):
        # 获取当前页
        text = ""
        try:
            page = doc.load_page(page_num)
            imgs = page.get_images()
            if len(imgs) == 1:
                # 可能整页就一个图片
                for img_index, img in enumerate(imgs):
                    print(
                        f"处理page:{page_num}/{len(doc)} 中的图片{(img_index+1)}/{len(imgs)}")
                    xref = img[0]
                    base_img = doc.extract_image(xref)
                    image_bytes = base_img['image']
                    image_ext = base_img['ext']
                    # 使用 tempfile 创建临时文件
                    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{image_ext}") as temp_file:
                        temp_file.write(image_bytes)
                        temp_file.flush()
                    temp_file_path = temp_file.name
                    # print(f"temp_file_path:{temp_file_path}")
                    text += read_img(temp_file_path)+"\n"
                    os.remove(temp_file_path)
                    # print(text[-100:])
            # 提取当前页的文字
            page_text = page.get_text("text")
            # 将提取的文字添加到总文字中
            text += page_text
            page = outdoc.new_page()
            text_rect = pymupdf.Rect(50, 50, 545, 800)  # 定义文本框区域（左、下、右、上）
            page.insert_textbox(
                text_rect,
                text,
                fontsize=14,
                align=0,  # 对齐方式：0=左对齐，1=居中，2=右对齐
                fontname="wqyyh",
                fontfile="fonts/"+config["def_font"]
            )

        except Exception as pdferr:
            print(pdferr)
        # 关闭文档
    doc.close()
    outdoc.save(output_file)
    outdoc.close()
    return text


if __name__ == "__main__":
    args = os.sys.argv
    if len(args) < 3:
        print("Usage: python pdfocrapp.py <input_file> <output_file>")
        exit(1)

    input_file = args[1]
    output_file = args[2]
    extract_text_from_pdf(input_file, output_file, 10)
    print(f"输入pdf已完成转换并保存在:{output_file}")
